'''
Brief:
  cnn训练部分代码
Requirement:
  完善训练主循环
'''

from model import *
from dataset import *
from opts import parser

#from torch.utils.tensorboard import SummaryWriter
from torch.nn.init import constant_, xavier_uniform_
import torch
import torchvision
import numpy as np 
import os
import shutil

def main():
    global args, best_prec1
    best_prec1 = 0
    #用于解析命令行参数和最高的acc来保存
    args = parser.parse_args()
    args_dict = args.__dict__
    print("------------------------------------")
    print("Configurations:")
    for key in args_dict.keys():
        print("- {}: {}".format(key, args_dict[key]))
    print("------------------------------------")
    #args:
    #  args.num_classes(int):命令行传入的分类个数
    #  args.train_list:训练集
    #  args.valid_list:交叉验证集
    #  args.batch_size(int):
    #  args.epochs(int):训练代数
    #  args.lr:学习率
    #  args.pretrained_model(str):预训练模型路径
    #  args.start_epoch:开始训练代数
    #  args.logdir:checkpoint保存路径
    #  args.eval_freq:评价频率,每隔多少代进行一次交叉验证并保存
    #device = "cpu"
    device = "cuda:0" if torch.cuda.is_available() else "cpu"

    #--------------初始化模型 -------------------------
    #-----------需要自行实现部分--------------------------
    #变量：
    #model(nn.Module)：自己实现的数字分类模型的实例
    #model_dict():模型的参数
    model = ClassifierNet(args.num_classes).to(device)
    model_dict = model.state_dict()

    #-----------初始化参数------------------
    if args.pretrained_model:
        if os.path.isfile(args.pretrained_model):
            print(("=> loading checkpoint '{}'".format(args.pretrained_model)))
            checkpoint = torch.load(args.pretrained_model)
            #checkpoint是个压缩文件，不仅包括权重，因此需要单独写一个转换的函数
            if "lr" not in checkpoint.keys():
                args.lr = input("No 'lr' attribute found in resume model, please input the 'lr' manually: ")
                args.lr = float(args.lr)
            else:
                args.lr = checkpoint['lr']
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']
            model.load_state_dict(checkpoint['state_dict'])
        else:
            print(("=> no checkpoint found at '{}'".format(args.pretrained_model)))
    else:
        #如果不存在预训练模型，那么需要对权重进行初始化
        new_state_dict = {}
        un_init_dict_keys = [k for k in model_dict.keys() if k not in new_state_dict]
        print("un_init_dict_keys: ", un_init_dict_keys)
        print("\n------------------------------------")

        for k in un_init_dict_keys:
            new_state_dict[k] = torch.DoubleTensor(model_dict[k].size()).zero_()
            if 'weight' in k:
                if 'bn' in k:
                    print("{} init as: 1".format(k))
                    constant_(new_state_dict[k], 1)
                else:
                    print("{} init as: xavier".format(k))
                    xavier_uniform_(new_state_dict[k])
            elif 'bias' in k:
                    print("{} init as: 0".format(k))
                    constant_(new_state_dict[k], 0)

        print("------------------------------------")
        model.load_state_dict(new_state_dict)
    
    logger = None
    #---------------------------需要自己实现的训练部分--------------------------#
    #-------变量------#
    #train_loader(torch.utils.data.DataLoader):训练的DataLoader
    #val_loader(torch.utils.data.DataLoader):交叉验证的DataLoader
    #criterion():loss function，可以自定义也可以使用torch现成的
    #optimizer():优化器，可以随便选，这里给的是Adam优化器
    train_loader = torch.utils.data.DataLoader(Dataset(args.train_list),batch_size=args.batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(Dataset(args.valid_list),batch_size=args.batch_size,shuffle=True)

    criterion = torch.nn.CrossEntropyLoss().to(device)#损失函数，交叉熵包括了softmax+log+nllloss
    optimizer = torch.optim.Adam(model.parameters(),args.lr)
    #--------------训练主循环----------#
    for epoch in range(args.start_epoch, args.epochs):
        train(train_loader,model,criterion,optimizer,epoch,device)

        # evaluate on validation set
        if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
            prec1 = validate(val_loader, model,criterion,epoch,device)

            # remember best prec@1 and save checkpoint
            is_best = prec1 > best_prec1
            best_prec1 = max(prec1, best_prec1)
            save_checkpoint({
                'epoch': epoch + 1,
                'state_dict': model.state_dict(),
                'best_prec1': best_prec1,
                'lr': optimizer.param_groups[-1]['lr'],
            }, is_best)
        
def train(train_loader,model,criterion, optimizer,epoch,device,logger=None):
    '''
    Brief:
      单个epoch的训练过程
    Args:
      train_loader(torch.utils.data.DataLoader):训练集的DataLoader
      model(torch.nn.Module):传入为模型实例
      criterion():损失函数
      optimizer():优化器
      epoch(int):当前轮数
      device(str):"cpu"或者"cuda:0"
      logger():日志记录，主要是用于Tensorboard可视化，这里删掉了(因为需要额外安装tensorflow)
        如果有兴趣可以用一下
    Return:
      无返回值
    Requirment:
      
    '''
    #包含某个评价标准的avg和当前值
    losses = AverageMeter()
    top1 = AverageMeter()


    torch.set_grad_enabled(True)
    model.train()

    loss_summ = 0

    #----------------加载dataloader-----------------#
    for i, (data,label) in enumerate(train_loader):
        #忽略最后一个batch
        if i == len(train_loader)-1:
            break

        #变为相关device的tensor
        input_var = data.to(device)
        target_var = label.to(device)

        output = model.forward(input_var)

        loss = criterion(output, target_var)
        loss_summ += loss
        #loss——sum用于计算平均loss
        #计算topk准确度
        prec1 = accuracy(output.data, label, topk=(1,))
        losses.update(loss_summ.item(), data.size(0))
        top1.update(prec1[0], data.size(0))

        #backprop
        loss.backward()

        optimizer.step()
        optimizer.zero_grad()
        loss_summ = 0
        #if i % args.print_freq == 0:
        print(('Epoch: [{0}][{1}/{2}], lr: {lr:.7f}\t'
                  'Loss {loss.val:.3f} ({loss.avg:.3f})\t'
                  'Prec@1 {top1.val:.2f} ({top1.avg:.2f})\t'.format(
                   epoch, i, len(train_loader), loss=losses, top1=top1, lr=optimizer.param_groups[-1]['lr'])))

def validate(val_loader, model, criterion,epoch,device,logger=None):
    losses = AverageMeter()
    top1 = AverageMeter()

    # In PyTorch 0.4, "volatile=True" is deprecated.
    torch.set_grad_enabled(False)

    # switch to evaluate mode
    model.eval()

    for i, (data,label) in enumerate(val_loader):
        # discard final batch
        if i == len(val_loader)-1:
            break
        input_var = data.to(device)
        target_var = label.to(device)

        # compute output
        output = model.forward(input_var)
        loss = criterion(output, target_var)

        # measure accuracy and record loss
        prec1 = accuracy(output.data, label, topk=(1,))

        losses.update(loss.item(), data.size(0))
        top1.update(prec1[0], data.size(0))

        if i % args.print_freq == 0:
            print(('Test: [{0}/{1}]\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'.format(
                   i, len(val_loader), loss=losses,
                   top1=top1)))

    print(('Testing Results: Prec@1 {top1.avg:.3f} Loss {loss.avg:.5f}'
          .format(top1=top1, loss=losses)))

    return top1.avg

def save_checkpoint(state:dict, is_best:bool, filename='checkpoint.pth.tar'):
    """
    Brief:
      以字典的形式压缩并保存模型和参考量
    Args:
      state(dict):传入字典
      is_best(bool):是否是最优的模型，如果是最优模型，那么会单独保存一次
      filename(str):不传入也行
    Example:
        save_checkpoint({'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'best_prec1': best_prec1,
            'lr': optimizer.param_groups[-1]['lr']}, is_best)
    """
    filename = '_'.join((args.logdir+'/checkpoints/',"epoch", str(state['epoch']), filename))
    torch.save(state, filename)
    if is_best:
        best_name = args.logdir+'/checkpoints/'+'model_best.pth.tar'
        shutil.copyfile(filename, best_name)

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        '''
        Brief:
          更新参数
        Args:
          val:loss/prec
          n(int):batch_size
        '''
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


if __name__ == '__main__':
    main()
